#include <clang/AST/ASTContext.h>
#include <clang/AST/RecursiveASTVisitor.h>
#include <clang/Frontend/FrontendActions.h>
#include <clang/Lex/Lexer.h>
#include <clang/Tooling/CommonOptionsParser.h>
#include <clang/Tooling/Tooling.h>
#include <clang/ASTMatchers/ASTMatchFinder.h>
#include <clang/ASTMatchers/ASTMatchers.h>
#include <llvm/Support/CommandLine.h>
#include <llvm/Support/raw_ostream.h>

using namespace clang;
using namespace clang::tooling;
using namespace clang::ast_matchers;

namespace {

llvm::cl::OptionCategory ToolCategory("stage08-tooling");

class ConditionalPrinter : public MatchFinder::MatchCallback {
public:
  void run(const MatchFinder::MatchResult &Result) override {
    const IfStmt *ifStmt = Result.Nodes.getNodeAs<IfStmt>("ifStmt");
    if (!ifStmt)
      return;

    auto &SM = *Result.SourceManager;
    auto loc = ifStmt->getIfLoc();
    if (loc.isInvalid())
      return;

    llvm::outs() << "[IfStmt] Found at "
                 << SM.getFilename(loc) << ":"
                 << SM.getSpellingLineNumber(loc) << "\n";
  }
};

class LoopTodoChecker : public MatchFinder::MatchCallback {
public:
  void run(const MatchFinder::MatchResult &Result) override {
    const ForStmt *forStmt = Result.Nodes.getNodeAs<ForStmt>("loop");
    if (!forStmt)
      return;

    const SourceManager &SM = *Result.SourceManager;
    SourceLocation loc = forStmt->getForLoc();
    if (loc.isInvalid())
      return;

    // Look for TODO comments near the loop (very naive check).
    SourceRange sourceRange(loc, forStmt->getEndLoc());
    StringRef text =
        Lexer::getSourceText(CharSourceRange::getTokenRange(sourceRange), SM,
                             Result.Context->getLangOpts());
    if (text.contains("TODO")) {
      llvm::outs() << "[Loop] TODO found near loop at "
                   << SM.getFilename(loc) << ":"
                   << SM.getSpellingLineNumber(loc) << "\n";
    }
  }
};

class ComputeCallChecker : public MatchFinder::MatchCallback {
public:
  void run(const MatchFinder::MatchResult &Result) override {
    const CallExpr *call = Result.Nodes.getNodeAs<CallExpr>("computeCall");
    if (!call)
      return;
    const FunctionDecl *callee = call->getDirectCallee();
    if (!callee)
      return;

    const SourceManager &SM = *Result.SourceManager;
    SourceLocation loc = call->getBeginLoc();
    llvm::outs() << "[Call] compute(...) used at "
                 << SM.getFilename(loc) << ":"
                 << SM.getSpellingLineNumber(loc) << "\n";
  }
};

} // namespace

int main(int argc, const char **argv) {
  auto ExpectedParser = CommonOptionsParser::create(argc, argv, ToolCategory);
  if (!ExpectedParser) {
    llvm::errs() << ExpectedParser.takeError();
    return 1;
  }

  CommonOptionsParser &OptionsParser = ExpectedParser.get();
  ClangTool Tool(OptionsParser.getCompilations(),
                 OptionsParser.getSourcePathList());

  ConditionalPrinter ifPrinter;
  LoopTodoChecker loopChecker;
  ComputeCallChecker computeChecker;

  MatchFinder finder;
  finder.addMatcher(ifStmt().bind("ifStmt"), &ifPrinter);
  finder.addMatcher(forStmt().bind("loop"), &loopChecker);
  finder.addMatcher(
      callExpr(callee(functionDecl(hasName("compute")))).bind("computeCall"),
      &computeChecker);

  return Tool.run(newFrontendActionFactory(&finder).get());
}
